﻿using DataProviderPluginContract;
using Liveforensics.Documents;
using Microsoft.Win32.SafeHandles;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Management;


namespace Liveforensics.RawProvider
{
    class SectorBuffer
    {
        const int BUFFER_SIZE = 65536;
        byte[] buffer;
        public byte[] Buffer
        {
            get { return buffer; }
            set { buffer = value; }
        }
        string diskName;
        public string DiskName
        {
            get { return diskName; }
            set { diskName = value; }
        }
        UInt64 startSector;
        public UInt64 StartSector
        {
            get { return startSector; }
            set { startSector = value; }
        }
        Int32 bufferSize;
        public Int32 BufferSize
        {
            get { return Buffer.Length; }
            private set { bufferSize = value; }
        }
        bool active;
        public bool Active
        {
            get { return active; }
            set { active = value; }
        }
        UInt64 diskSize;
        public UInt64 DiskSize
        {
            get { return diskSize; }
            set { diskSize = value; }
        }
        DiskInfo diskInformation;
        public DiskInfo DiskInformation
        {
            get { return diskInformation; }
            set { diskInformation = value; }
        }
        public SectorBuffer()
        {
            Buffer = new byte[BUFFER_SIZE];
        }
        public byte[] ReadSectors(ulong start, uint count)
        {
            try
            {
                if (!Active)
                    return null;
                if (start < StartSector)
                    return null;
                uint totalSectorsInBuffer = BUFFER_SIZE / (uint)DiskInformation.BytesPerSector;
                ulong lastSector = totalSectorsInBuffer + StartSector;
                if (start + count - 1 > lastSector)
                    return null;
                ulong bytePointer = (start - StartSector) * DiskInformation.BytesPerSector;
                uint bufferLength = count * DiskInformation.BytesPerSector;
                byte[] outputBuffer = new byte[bufferLength];
                for (uint i = 0; i < bufferLength; i++)
                {
                    outputBuffer[i] = Buffer[i + bytePointer];
                }
                return outputBuffer;
            }
            catch
            {
                return null;
            }
        }
    }
    public class RawProvider : IDataProviderPlugin
    {
        #region IPlugin Members
        List<DiskInfo> _physicalDisks;
        List<SectorBuffer> _sectorBuffers;

        public RawProvider()
        {
            _physicalDisks = new List<DiskInfo>();
            _sectorBuffers = new List<SectorBuffer>();
        }
        public bool Load(string targetDisk)
        {
            return false;
        }
        public Guid identifier
        {
            get
            {
                return new Guid("5D16A0E5-1E7A-4E71-83BF-17815BD0879C");
            }
            
        }
        public string Name
        {
            get { return "Local Machine";}
        }
        public string ImageName
        {
            get
            {
                return "";
            }
        }
        public List<string> FileExtensions
        {
            get
            {
                // first entry is the default extension e.g. ".png"
                // then a list of filter expressions
                // eg "JPEG Files (*.jpeg)|*.jpeg|PNG Files (*.png)|*.png|JPG Files (*.jpg)|*.jpg|GIF Files (*.gif)|*.gif"
                List<string> extensions = new List<string>();
                extensions.Add("*.vmdk");
                extensions.Add("VMDK Files (*.vmdk)");
                extensions.Add("*.vmdk");
                return extensions;
            }
        }
        public bool needsAdmin
        {
            get
            {
                return true;
            }
        }
        public void Do()
        {
            Debug.WriteLine("Hello from the Raw Provider Plugin");
        }
        public long GetDiskGeometry(DiskInfo di)
        {
            SafeFileHandle handle = Disk.OpenFileHandle(di.DeviceID);
            NativeMethods.DiskGeometry dg = Disk.GetDiskGeometry(handle);

            return dg.DiskSize;
        }
        public List<DiskInfo> GetPhysicalDisks()
        {
            if (_physicalDisks == null)
                _physicalDisks = new List<DiskInfo>();
            else
                _physicalDisks.Clear();
            WqlObjectQuery q = new WqlObjectQuery("SELECT * FROM Win32_DiskDrive");
            ManagementObjectSearcher res = new ManagementObjectSearcher(q);
            foreach (ManagementObject o in res.Get())
            {
                try
                {
                    DiskInfo di = new DiskInfo();
                    di.Identifier = Guid.NewGuid();
                    try
                    {
                        di.Name = o["Name"].ToString();
                        if (di.Name.StartsWith(@"\\.\"))
                            di.Name = di.Name.Substring(4);
                    }
                    catch { di.Name = "empty"; }
                    try { di.Manufacturer = o["Manufacturer"].ToString(); }
                    catch { di.Manufacturer = "empty"; }
                    try { di.MediaType = o["MediaType"].ToString(); }
                    catch { di.MediaType = "empty"; }
                    try { di.Model = o["Model"].ToString(); }
                    catch { di.Model = "empty"; }
                    try { di.SerialNumber = o["SerialNumber"].ToString(); }
                    catch { di.SerialNumber = "empty"; }
                    try { di.PartitonCount = (Int32)o["Partitions"]; }
                    catch { di.PartitonCount = 0; }
                    try { di.ReportedSize = (UInt64)o["Size"]; }
                    catch { di.ReportedSize = 0; }
                    try { di.Caption = o["Caption"].ToString(); }
                    catch { di.Caption = "empty"; }
                    try { di.Description = o["Description"].ToString(); }
                    catch { di.Description = "empty"; }
                    try { di.DeviceID = o["DeviceID"].ToString(); }
                    catch { di.DeviceID = "empty"; }
                    try { di.BytesPerSector = (uint)o["BytesPerSector"];} 
                    catch 
                    {
                        di.BytesPerSector = Disk.GetBytesPerSector(di.DeviceID);
                    }
                    ulong tempSize = (ulong)GetDiskGeometry(di);
                    if (di.ReportedSize == 0)
                    {
                        try
                        {
                            di.ReportedSize = (ulong)GetDiskGeometry(di);
                        }
                        catch
                        {
                            di.ReportedSize = 0;
                        }
                    }
                    //di.ActualSize = GetRealDiskSize(di);
                    _physicalDisks.Add(di);
                    SectorBuffer sb = new SectorBuffer();
                    sb.DiskInformation = di;
                    sb.DiskSize = di.ReportedSize;
                    sb.Active = false;
                    _sectorBuffers.Add(sb);
                }
                catch { }
            }


            return _physicalDisks;
        }
        private ulong GetRealDiskSize(DiskInfo d)
        {
            ulong startingSector = d.ReportedSize / d.BytesPerSector;
            byte[] result = null;
            ulong gap = 1024;
            bool trigger = false;
            while (true)
            {
                Debug.Write("Reading: " + (startingSector + gap).ToString());
                result = ReadSectorNoCache(d, startingSector + gap);
                if (result != null)
                {
                    Debug.WriteLine(" Success");
                    if (gap > 128)
                        ReadSector(d, startingSector + gap);
                    startingSector += gap;
                }
                else
                {
                    Debug.WriteLine(" Failed");
                    trigger = true;
                }
                if (trigger)
                    gap /= 2;
                if (gap == 0)
                    break;
            }
            return startingSector * d.BytesPerSector;
        }
        private byte[] ReadSectorNoCache(DiskInfo di, ulong startSector, uint sectorCount = 1)
        {
            SectorBuffer activeSectorBuffer = new SectorBuffer();
            activeSectorBuffer.DiskInformation = di;
            if (di != null)
            {
                byte[] test = ReadFromCache(di, startSector, sectorCount);
                if (test != null)
                    return test;
                //long realSize = di.BytesPerSector * sectorCount;
                try
                {
                    uint bps = activeSectorBuffer.DiskInformation.BytesPerSector;
                    long requestedSize = sectorCount * bps;
                    SafeFileHandle handle = Disk.OpenFileHandle(di.DeviceID);
                    byte[] readBuffer = Disk.ReadSectors(handle, startSector, sectorCount);
                    return readBuffer;
                }
                catch { return null; }     
            }
            return null;
        }
        private byte[] ReadFromCache(DiskInfo di, ulong startSector, uint sectorCount = 1)
        {
            if (di != null)
            {
                foreach (SectorBuffer sb in _sectorBuffers)
                {
                    if (di == sb.DiskInformation)
                    {
                        SectorBuffer activeSectorBuffer = sb;
                        byte[] buffer = sb.ReadSectors(startSector, sectorCount);
                        if (buffer != null)
                            return buffer;
                        break;
                    }
                }
            }
            return null;
        }
        public byte[] ReadSector(DiskInfo di, ulong startSector, uint sectorCount = 1)
        {
            SectorBuffer activeSectorBuffer = new SectorBuffer();
            activeSectorBuffer.DiskInformation = di;
            byte[] readBuffer = ReadFromCache(di, startSector, sectorCount);
            if (readBuffer != null)
                return readBuffer;
            if (di != null)
            {                
                try
                {
                    uint bps = activeSectorBuffer.DiskInformation.BytesPerSector;
                    long requestedSize = sectorCount * bps;
                    SafeFileHandle handle = Disk.OpenFileHandle(di.DeviceID);
                    if (requestedSize > activeSectorBuffer.BufferSize)
                        return Disk.ReadSectors(handle, startSector, sectorCount);

                    uint maximumEndSector = (uint)(di.ActualSize / bps) - ((uint)activeSectorBuffer.BufferSize / bps);
                    readBuffer = null;
                    if (startSector > maximumEndSector)
                    {
                        readBuffer = Disk.ReadSectors(handle, maximumEndSector, (uint)activeSectorBuffer.BufferSize / bps);
                        activeSectorBuffer.StartSector = maximumEndSector;
                    }
                    else
                    {
                        readBuffer = Disk.ReadSectors(handle, startSector, (uint)activeSectorBuffer.BufferSize / bps);
                        activeSectorBuffer.StartSector = startSector;
                    }                    
                    activeSectorBuffer.Active = true;
                    Debug.WriteLine("Cached: " + activeSectorBuffer.StartSector.ToString() + " Length: " + (activeSectorBuffer.BufferSize / bps).ToString());
                    for (int i = 0; i < activeSectorBuffer.BufferSize; i++)
                        activeSectorBuffer.Buffer[i] = readBuffer[i];
                    byte[] buffer = activeSectorBuffer.ReadSectors(startSector, sectorCount);
                    if (buffer != null)
                        return buffer;
                }
                catch { return null; }                
            }
            return null;
        }
        

        #endregion
    }    
}
